import numpy as np

from sklearn import metrics





def get_confusion_matrix(real_mask,pred_mask,n_class=5):
    '''
   适用于部分图像没有某个类别标签的情况。例如2D肿瘤分割，很多图片没有肿瘤。
   这时候要先计算所有图像的混淆矩阵，再相加，然后利用总的混淆矩阵计算各项评估指标
   '''
    y_true = real_mask.reshape(1, -1).squeeze()
    y_pred = pred_mask.reshape(1, -1).squeeze()
    confusion_matrix = metrics.confusion_matrix(y_true, y_pred,labels = [i for i in range(n_class)])
    return np.array(confusion_matrix)

class Meterics:
    def __init__(self,n_class):
        self.n_class=n_class
        self.confusion_matrix_total=np.zeros((n_class, n_class))
    def _get_confusion_matrix(self,real_mask, pred_mask):
        '''
       计算混淆矩阵
       '''
        y_true = real_mask.reshape(1, -1).squeeze()
        y_pred = pred_mask.reshape(1, -1).squeeze()
        confusion_matrix = metrics.confusion_matrix(y_true, y_pred, labels=[i for i in range(self.n_class)])
        return np.array(confusion_matrix)
    def add_batch(self,real_mask, pred_mask):
        '''
        将混淆矩阵累加
        '''
        confusion=self._get_confusion_matrix(real_mask,pred_mask)
        self.confusion_matrix_total+=confusion
    def get_iou(self):
        return np.diag(self.confusion_matrix_total) / (self.confusion_matrix_total.sum(axis=1)
                + self.confusion_matrix_total.sum(axis=0) - np.diag(self.confusion_matrix_total))
    def get_dice(self):
        return (2 * np.diag(self.confusion_matrix_total)) / ((self.confusion_matrix_total.sum(axis=1)
                + self.confusion_matrix_total.sum(axis=0)))
    def get_precision(self):
        return np.diag(self.confusion_matrix_total) / self.confusion_matrix_total.sum(axis=0)
    def get_recall(self):
        return np.diag(self.confusion_matrix_total) / self.confusion_matrix_total.sum(axis=1)
    def get_f1(self):
        #f1和dice其实是一样的
        return self.get_dice()



if __name__ == '__main__':
    y_true=np.array([0,1,2,0,1,0,1,0])#true
    y_pre=np.array( [1,1,2,0,0,2,0,1])#pre
    '''
    第i行、第j列表示：真实为i的，预测为j的
    [[1 2 1]
     [2 1 0]
     [0 0 1]]
     0：inter=1,union=6
     1: inter=1,union=5
     第i行的和代表真实为i，预测的所有结果
     第j列的和表示，预测结果为j的所有结果的和
     precesion=预测正确（i=j）/第j列的和
     recall=预测正确（i=j）/第i行的和
     '''
    # self.tp = confusion_matrix[0][0]
    # self.fn = confusion_matrix[0][1]
    # self.fp = confusion_matrix[1][0]
    # # self.tn = confusion_matrix[1][1]
    m=Meterics(n_class=2)
    result=get_confusion_matrix(y_true,y_pre,3)#第三个参数是类别数
    #self.tn, self.fp, self.fn, self.tp = confusion_matrix.ravel()  # tn, fp, fn, tp
    print(result)
    iu = np.diag(result) / (result.sum(axis=1) + result.sum(axis=0) - np.diag(result))
    a=result.sum(axis=0)#求每一列的和
    b=result.sum(axis=1)#求每一行的和
    dice=(2*np.diag(result))/((result.sum(axis=1) + result.sum(axis=0)))
    #预测为true的，有多少为true
    precision=np.diag(result) / result.sum(axis=0)
    recall=np.diag(result) / result.sum(axis=1)
    f1=(2*precision*recall)/(precision+recall)
    print(iu)
    print(dice)
    print(precision)
    print(recall)
    print(f1)
    exit()